import dgl
import torch

def get_total_graph(rn):
    """
    Build the total graph
    """
    src = []
    dst = []
    for rid in rn.valid_edge.keys():
        for eid in rn.edgeDict[rid]:
            if eid in rn.valid_edge.keys():
                src.append(rn.valid_edge_one[rid])
                dst.append(rn.valid_edge_one[eid])
        for eid in rn.edgeRevDict[rid]:
            if eid in rn.valid_edge.keys():
                src.append(rn.valid_edge_one[rid])
                dst.append(rn.valid_edge_one[eid])
    
    g = dgl.DGLGraph()
    g.add_nodes(rn.valid_edge_cnt_one)
    g.add_edges(src, dst)
    g = dgl.add_self_loop(g)
    
    return g

def get_sub_graphs(rn, max_deps=3):
    """
    Build sub graphs for each rid
    """
    subg = {}
    for rid in rn.valid_edge.keys():
        rset = set()
        rset.add(rid)
        for dep in range(max_deps):
            new_rset = set()
            for r in rset:
                for eid in rn.edgeDict[r]:
                    if eid in rn.valid_edge.keys():
                        new_rset.add(eid)
                for eid in rn.edgeRevDict[r]:
                    if eid in rn.valid_edge.keys():
                        new_rset.add(eid)
            rset = rset.union(new_rset)
        
        rset = list(rset)
        mset = [rn.valid_edge_one[rid] for rid in rset]
        rmap = {}
        for i, r in enumerate(rset):
            rmap[rn.valid_edge_one[r]] = i
        
        src = []
        dst = []
        for r in rset:
            for eid in rn.edgeDict[r]:
                if eid in rset:
                    src.append(rmap[rn.valid_edge_one[r]])
                    dst.append(rmap[rn.valid_edge_one[eid]])
            for eid in rn.edgeRevDict[r]:
                if eid in rset:
                    src.append(rmap[rn.valid_edge_one[r]])
                    dst.append(rmap[rn.valid_edge_one[eid]])
        
        g = dgl.DGLGraph()
        g.add_nodes(len(rset))
        g.add_edges(src, dst)
        g = dgl.add_self_loop(g)
        
        subg[rn.valid_edge_one[rid]] = g
    
    return subg

def empty_graph():
    """
    Build an empty graph
    """
    g = dgl.DGLGraph()
    g.add_nodes(1)
    g = dgl.add_self_loop(g)
    return g